import os
import os.path
import sys
import math
import argparse
import time
import random
import numpy as np
import utils
from torchvision.utils import save_image, make_grid
from collections import OrderedDict
from tqdm import tqdm

import torch
import torchvision
import torch.optim as optim

from datasets import create_dataloader, create_dataset
from models import create_network
from torch.autograd import Variable
from math import log10


EXP = "PPT"
TESTSET = "SceneText"
EPOCH = 120
TEST_MODE_GPU = True
FORWARD_NAME = "epoches/"+EXP+"/forward_net_epoch_"+str(EPOCH)+".pth"


if __name__ == "__main__":
    if not os.path.exists("results/"+TESTSET):
        os.makedirs("results/"+TESTSET)
    # create dataloader
    test_set = create_dataset("LR", "val", TESTSET)
    test_loader = create_dataloader(test_set, "test")

    # create model
    forward_net, _a1, _a2 = create_network()

    if TEST_MODE_GPU:
        forward_net.cuda()
        forward_net.load_state_dict(torch.load(FORWARD_NAME))
    else:
        forward_net.load_state_dict(
            (torch.load(FORWARD_NAME, map_location=lambda storage, loc: storage)))

    test_bar = tqdm(test_loader, desc="[converting, saving and calculating]")
    total_psnr = 0
    with torch.no_grad():
        index = 1
        for lr in test_bar:
            if TEST_MODE_GPU:
                lr = lr.cuda()
            b, c, h, w = lr.shape
            sr = forward_net(lr)
            torchvision.utils.save_image(
                sr, ("results/"+TESTSET+"/index_"+str(index)+".jpg"))
            index += 1
